// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <fmt/format.h>
#include <melon/benchmark.h>
#include <melon/conv.h>
#include <melon/init/init.h>
#include <string>

#include <pollux/functions/lib/re2_functions.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/vector/vector_fuzzer.h>

namespace kumo::pollux::functions::test {
namespace {

int regexMatch(int n, int blockSize, const char* functionName) {
  melon::BenchmarkSuspender kSuspender;
  FunctionBenchmarkBase benchmarkBase;

  VectorFuzzer::Options opts;
  opts.vectorSize = blockSize;
  auto vector = VectorFuzzer(opts, benchmarkBase.pool()).fuzzFlat(VARCHAR());
  const auto data = benchmarkBase.maker().row_vector({vector});

  exec::ExprSet expr = benchmarkBase.compileExpression(
      melon::to<std::string>(functionName, "(c0, '[^9]{3,5}')"), data->type());
  kSuspender.dismiss();
  for (int i = 0; i != n; ++i) {
    benchmarkBase.evaluate(expr, data);
  }
  return n * blockSize;
}

BENCHMARK_NAMED_PARAM_MULTI(regexMatch, bs1k, 1 << 10, "re2_match");
BENCHMARK_NAMED_PARAM_MULTI(regexMatch, bs10k, 10 << 10, "re2_match");
BENCHMARK_NAMED_PARAM_MULTI(regexMatch, bs100k, 100 << 10, "re2_match");

int regexSearch(int n, int blockSize, const char* functionName) {
  return regexMatch(n, blockSize, functionName);
}

BENCHMARK_NAMED_PARAM_MULTI(regexSearch, bs1k, 1 << 10, "re2_search");
BENCHMARK_NAMED_PARAM_MULTI(regexSearch, bs10k, 10 << 10, "re2_search");
BENCHMARK_NAMED_PARAM_MULTI(regexSearch, bs100k, 100 << 10, "re2_search");

int regexExtract(int n, int blockSize) {
  melon::BenchmarkSuspender kSuspender;
  FunctionBenchmarkBase benchmarkBase;

  VectorFuzzer::Options opts;
  opts.vectorSize = blockSize;
  auto vector = VectorFuzzer(opts, benchmarkBase.pool()).fuzzFlat(VARCHAR());
  const auto data = benchmarkBase.maker().row_vector({
      vector,
      benchmarkBase.maker().constant_vector<int32_t>({0}),
  });

  exec::ExprSet expr = benchmarkBase.compileExpression(
      "re2_extract(c0, '99[^9]', c1)", data->type());
  kSuspender.dismiss();
  for (int i = 0; i != n; ++i) {
    benchmarkBase.evaluate(expr, data);
  }
  return n * blockSize;
}

BENCHMARK_NAMED_PARAM_MULTI(regexExtract, bs100, 100);
BENCHMARK_NAMED_PARAM_MULTI(regexExtract, bs1k, 1 << 10);
BENCHMARK_NAMED_PARAM_MULTI(regexExtract, bs10k, 10 << 10);
BENCHMARK_NAMED_PARAM_MULTI(regexExtract, bs100k, 100 << 10);

} // namespace

std::shared_ptr<exec::VectorFunction> makeRegexExtract(
    const std::string& name,
    const std::vector<exec::VectorFunctionArg>& inputArgs,
    const core::QueryConfig& config) {
  return makeRe2Extract(name, inputArgs, config, /*emptyNoMatch=*/false);
}

void registerRe2Functions() {
  exec::registerStatefulVectorFunction(
      "re2_match", re2MatchSignatures(), makeRe2Match);
  exec::registerStatefulVectorFunction(
      "re2_search", re2SearchSignatures(), makeRe2Search);
  exec::registerStatefulVectorFunction(
      "re2_extract", re2ExtractSignatures(), makeRegexExtract);
}

} // namespace kumo::pollux::functions::test

int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};
  kumo::pollux::functions::test::registerRe2Functions();
  kumo::pollux::memory::MemoryManager::initialize({});
  melon::runBenchmarks();
  return 0;
}
